{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 176,
   "id": "1e6fdad0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader, ConcatDataset\n",
    "from torch.autograd import Variable\n",
    "import torchvision\n",
    "from torchvision import transforms, models\n",
    "import matplotlib.pyplot as plt\n",
    "import multiprocessing\n",
    "from sklearn.metrics import accuracy_score\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "id": "0a8c0b6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SpaceDataset(Dataset):\n",
    "    def __init__(self, path):\n",
    "        self.path = path\n",
    "        self.img_list = os.listdir(path + 'feature/')\n",
    "        print(self.img_list)\n",
    "        self.preprocess = transforms.Compose([transforms.ToTensor()])\n",
    "        score_file = open(path + \"score.txt\")\n",
    "        self.scores = score_file.read().splitlines()\n",
    "        score_file.close()\n",
    "        print(\"len(img_list)={0}\".format(len(self.img_list)))\n",
    "        print(\"len(scores)={0}\".format(len(self.scores)))\n",
    "        assert(len(self.img_list) == len(self.scores))\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        img_filename = self.img_list[index]\n",
    "        img_path = self.path + \"feature/\" + img_filename\n",
    "        img = Image.open(img_path)\n",
    "        img = self.preprocess(img.convert(\"RGB\"))\n",
    "        label_idx = int(img_filename[-8:-4:1])\n",
    "        label = torch.tensor([float(self.scores[label_idx])])\n",
    "        return img, label\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.img_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "id": "1212c806",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, dataset, optimizer, scheduler, loss, batch_size, epochs):\n",
    "    #model = model.cuda()\n",
    "    for epoch in range(epochs):\n",
    "        loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "        print(\"<<<<<<< enter epoch {0} >>>>>>>>>\".format(epoch+1))\n",
    "        for img, label in loader:\n",
    "            img = Variable(img, requires_grad=True)\n",
    "            label = Variable(label, requires_grad=True)\n",
    "            #img = Variable(img, requires_grad=True).cuda()\n",
    "            #label = Variable(label, requires_grad=True).cuda()\n",
    "            predict = model(img)\n",
    "            type(predict)\n",
    "            loss(predict, label).backward()\n",
    "            optimizer.step()\n",
    "            optimizer.zero_grad()\n",
    "            del img\n",
    "            print(\"<<<<<<< predict - target = {0} >>>>>>>>>\".format(abs(predict-label)[0][0]))\n",
    "        scheduler.step()\n",
    "    torch.save(model.state_dict(), './model.weight')\n",
    "    return \"training complete!\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "id": "165fefd3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['f0002.png', 'f0003.png', 'f0001.png', 'f0000.png', 'f0004.png', 'f0005.png', 'f0007.png', 'f0006.png', 'f0008.png', 'f0009.png']\n",
      "len(img_list)=10\n",
      "len(scores)=10\n"
     ]
    }
   ],
   "source": [
    "resnet101 = models.resnet101(pretrained = True)\n",
    "resnet101.fc = torch.nn.Linear(in_features=2048, out_features=1, bias=True)\n",
    "#resnet101.cuda()\n",
    "data = SpaceDataset(\"dataset/\")\n",
    "opt = torch.optim.Adam(resnet101.parameters(), lr=0.001)\n",
    "lr_adjuster = torch.optim.lr_scheduler.StepLR(opt, step_size=10, gamma=0.1)\n",
    "criterion = torch.nn.MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "id": "e9a35941",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<<<<<<< enter epoch 1 >>>>>>>>>\n",
      "<<<<<<< predict - target = 96.72421264648438 >>>>>>>>>\n",
      "<<<<<<< predict - target = 64.29517364501953 >>>>>>>>>\n",
      "<<<<<<< predict - target = 20.53515625 >>>>>>>>>\n",
      "<<<<<<< predict - target = 29.6707820892334 >>>>>>>>>\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-193-f8b610b41071>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresnet101\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mopt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlr_adjuster\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-185-b12ee441253d>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(model, dataset, optimizer, scheduler, loss, batch_size, epochs)\u001b[0m\n\u001b[1;32m     11\u001b[0m             \u001b[0mpredict\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m             \u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpredict\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 13\u001b[0;31m             \u001b[0mloss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpredict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     14\u001b[0m             \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m             \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    243\u001b[0m                 \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    244\u001b[0m                 inputs=inputs)\n\u001b[0;32m--> 245\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    246\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    247\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    143\u001b[0m         \u001b[0mretain_graph\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    144\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 145\u001b[0;31m     Variable._execution_engine.run_backward(\n\u001b[0m\u001b[1;32m    146\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    147\u001b[0m         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "train(resnet101, data, opt, lr_adjuster, loss, 1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c984a5d7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
